import copy
import torch
from utils import *
from exp08.comodel import ComModel
from exp08.strategy import logits_to_span_list_with_mark
from tqdm import tqdm


def test(
        args,
        model: ComModel,
        test_data_loader,
        three_goden_set,
        logger,
        is_test: bool = False,  # 如果是 train 和 dev 数据集，那么 is_test 就是 false
        model_dir='',
):
    model.eval()
    if is_test:  # 如果使用的是 test 数据集
        model.load_state_dict(torch.load(model_dir))
        pass
    aspect_set, opinion_set, triplets_set, pair_set, as_set = set(), set(), set(), set(), set()
    # aspect_set 是一个 set 其中元素是形如 batch_i-start_i-end_i 的字符串
    # pair_set 是 ["i-al-ar-ol-or"]
    # triplets_set 是 [“i-al-ar-ol-or-s"]
    # as_set 是 [”i-al-ar-s“]
    multi_set, single_set = set(), set()
    # multi_set 和 single_set 都是 [“i-al-ar-ol-or-s"]

    multi_aspect_id = three_goden_set[2][3]  # original_data.multi_aspect_id 记录的是所有多标签的 id，从 0 开始

    def loop_calculate(input_plus: dict, token2id_list: list, from_first: bool = True):
        has_aspect: bool = True
        tmp_input_clone = copy.deepcopy(input_plus)
        count_extracted = 0
        res = []
        while has_aspect:
            model_res = model(tmp_input_clone, args)
            # model_res 当中包括的属性有 as_p, ae_p, as_p2, ae_p2, is_on_logits, os_p, oe_p, os_p2, oe_p2, s_logits, s_logits2
            as_p = model_res["as_p"] if from_first else model_res["as_p2"]
            ae_p = model_res["ae_p"] if from_first else model_res["ae_p2"]
            os_p = model_res["os_p"] if from_first else model_res["os_p2"]
            oe_p = model_res["oe_p"] if from_first else model_res["oe_p2"]
            s_logits = model_res['s_logits'] if from_first else model_res["s_logits2"]
            s = torch.argmax(s_logits, dim=1).squeeze().item()
            is_on_logits = model_res['is_on_logits']  # (batch_size, 2)
            has_aspect = torch.argmax(is_on_logits, dim=1).squeeze().item() == 1
            scored_aspect_span = logits_to_span_list_with_mark(as_p.squeeze(), ae_p.squeeze(), token2id_list)
            scored_opinion_span = logits_to_span_list_with_mark(os_p.squeeze(), oe_p.squeeze(), token2id_list, max_num=3)
            if not has_aspect or len(scored_aspect_span) == 0 or len(scored_opinion_span) == 0:
                break
                pass
            # 此时应该是能提取出 opinion 的，此时应该保存这个 opinion 的结果
            res.append([scored_aspect_span[0], scored_opinion_span, s])
            # 更新 mask
            inputs_plus['attention_mask'][0][args.sen_pre_len + scored_aspect_span[0][0]:args.sen_pre_len + scored_aspect_span[0][1] + 1] = 0
            count_extracted += 1
            if count_extracted > 6:
                break
                pass
            pass
        return res

    n_batch_called = 0
    with torch.no_grad():
        for batch_i, test_batch_data in tqdm(enumerate(test_data_loader), desc="testing" if is_test else "validating", delay=0.5):
            n_batch_called += 1
            is_on, count = True, 0  # count 代表的是这一个 test 的样本中执行了多少轮
            inputs_plus = test_batch_data['inputs_plus_for_test']
            token2id = test_batch_data["test_instances"][0].meta_instance.token2id

            res1 = loop_calculate(test_batch_data, token2id)
            res2 = loop_calculate(test_batch_data, token2id, from_first=False)

            pass
        pass

    a_p, a_r, a_f = score_set(aspect_set, three_goden_set[0])  # original_data.get_all_aspect_set(),
    o_p, o_r, o_f = score_set(opinion_set, three_goden_set[1])  # original_data.get_all_opinion_set(),
    t_p, t_r, t_f = score_set(triplets_set, three_goden_set[2][0])  # triplets_set 是 [“i-al-ar-ol-or-s"] 第二个是 original_data.all_triplets,
    mt_p, mt_r, mt_f = score_set(multi_set, three_goden_set[2][1])  # original_data.all_multi_triplets,
    st_p, st_r, st_f = score_set(single_set, three_goden_set[2][2])  # original_data.all_single_triplets,

    logger.info(get_dict_str_in_lines({
        "aspect.p r f1": " ".join(['%.5f' % a_p, '%.5f' % a_r, '%.5f' % a_f]),
        "opinion.p r f1": " ".join(['%.5f' % o_p, '%.5f' % o_r, '%.5f' % o_f]),
        "multi_t.p r f1": " ".join(['%.5f' % mt_p, '%.5f' % mt_r, '%.5f' % mt_f]),
        "single_t.p r f1": " ".join(['%.5f' % st_p, '%.5f' % st_r, '%.5f' % st_f]),
        "all_t.p r f1": " ".join(['%.5f' % t_p, '%.5f' % t_r, '%.5f' % t_f]),
    }))

    return mt_f, st_f, t_f
